#include "vector"
#include <opencv2/opencv.hpp>
#include <boost/format.hpp>
#include "Eigen/Dense"
#include <opencv2/core/eigen.hpp>
#include <opencv2/features2d/features2d.hpp>
#include <ceres/ceres.h>
#include <sophus/se3.hpp>
#include <iostream>

#define IMG_SIZE 2


struct ProjCost
{
  ProjCost(Eigen::Vector2d pos2d, Eigen::Vector3d pos3d, Eigen::Matrix3d K, Sophus::SE3d T21)
  : _pos2d(pos2d), _pos3d(pos3d), _K(_K), _T21(T21), _K_inv(K.inverse()) {};

  template<typename T>
  bool operator()(
    const T *const abc, T *residual) const {

    Eigen::Matrix<double, 6, 1> tmp;
    tmp << T(abc[0]), 0, 0, 0, 0, 0;
    Sophus::SE3d T_ = Sophus::SE3d::exp(tmp);
    const double scale = 1;
    Eigen::Vector3d pos_pixel = _K * (T_ * _T21 * (scale * _K_inv * _pos3d));
    pos_pixel /= pos_pixel[2];
    residual[0] = T((_pos2d - pos_pixel.head<2>()).norm());

    return true;
  }

  Eigen::Vector2d _pos2d;
  Eigen::Vector3d _pos3d;
  Eigen::Matrix3d _K;
  Sophus::SE3d _T21;
  Eigen::Matrix3d _K_inv;

};

int main()
{

    double fx = 718.856, fy = 718.856, cx = 607.1928, cy = 185.2157;

    std::vector<std::vector<int>> bbox_list = {
        {535, 169, 560, 229},
        {535, 170, 560, 235}
    };

    std::vector<Eigen::Matrix4d, Eigen::aligned_allocator<double>> T_list;
    T_list.resize(2);
    T_list[0] << 
    0.999991,  0.00248079,  0.00343383, -0.00373836,
    -0.00248834,    0.999994, 0.00219452,  0.00304478,
    -0.00342836, -0.00220304,    0.999992,   -0.732342,
    0,           0,           0,           1;
    
    T_list[1] << 
    0.999971,  0.000699793,   0.00759075,  -0.00252071,
    -0.000728066,     0.999993,    0.0037225,   0.00402008,
    -0.00758809,  -0.00372792,     0.999964,      -1.4814,
    0,            0,          0,           1;
    
    boost::format fmt_others("../%06d.png");    // other files
    std::vector<cv::Mat> img_list;
    img_list.resize(IMG_SIZE);
    std::vector<std::vector<cv::KeyPoint>> kps_list;
    kps_list.resize(IMG_SIZE);
    std::vector<cv::Mat> dscp_list;
    dscp_list.resize(IMG_SIZE);
    cv::Ptr<cv::Feature2D> detector = cv::ORB::create(10);
    cv::Ptr<cv::DescriptorExtractor> descriptor = cv::ORB::create();
    cv::Ptr<cv::DescriptorMatcher> matcher = cv::DescriptorMatcher::create("BruteForce-Hamming");


    for (int i = 1; i < 3; i++) {
        cv::Mat img = cv::imread((fmt_others % i).str(), 0);
        int j = i - 1;
        img_list[j] = img;
        
        // detect the key points
        cv::Point2i p1(bbox_list[j][0], bbox_list[j][1]);
        cv::Point2i p2(bbox_list[j][2], bbox_list[j][3]);

        cv::Rect rect{p1, p2};

        cv::Mat mask = cv::Mat::zeros(img.rows, img.cols, CV_8UC1);
        cv::rectangle(mask, p1, p2, (255), -1);
        cv::Mat roi(img, rect);
        detector->detect(img, kps_list[j], mask);
        std::cout << "N keypoints: " << kps_list[j].size() << std::endl;
        descriptor->compute(img, kps_list[j], dscp_list[j]);

        // take the T matrix
        Eigen::Matrix4d T = T_list[j];

        // plot
        cv::Mat color_bg;
        cv::cvtColor(img, color_bg, CV_GRAY2BGR);
        cv::rectangle(color_bg, p1, p2, {0, 255, 0}, 1, 1);
        cv::drawKeypoints(color_bg, kps_list[j], color_bg, cv::Scalar::all(-1), cv::DrawMatchesFlags::DEFAULT);
        // cv::imshow("key points", color_bg);
        // cv::waitKey();
    }
    // calculate the descriptor
    std::vector<cv::DMatch> matches;
    matcher->match(dscp_list[0], dscp_list[1], matches);

    //-- 第五步:绘制匹配结果
    // cv::Mat img_match;
    // cv::drawMatches(img_list[0], kps_list[0], img_list[1], kps_list[1], matches, img_match);
    // cv::imshow("all matches", img_match);
    // cv::waitKey();

    // begin to estimate the position
    double abc[7] = {0, 0, 0, 0, 0, 0, 1};
    
    // construct the T21
    const Eigen::Matrix4d T1W = T_list[0];
    const Eigen::Matrix4d T2W = T_list[1];
    const Eigen::Matrix4d T21_hat = T2W * T1W.inverse();

    Eigen::JacobiSVD<Eigen::MatrixXd> svd(T21_hat.block(0, 0, 3, 3), Eigen::ComputeThinU | Eigen::ComputeThinV);
    auto U = svd.matrixU();
    auto V = svd.matrixV(); 
    Sophus::SE3d T21(U * V, T21_hat.block(0, 3, 3, 1));

    // the K matrix
    Eigen::Matrix3d K;
    K <<
            fx, 0, cx,
            0, fy, cy,
            0, 0, 1;

    // 构建最小二乘问题
    ceres::Problem problem;
    for (int i = 0; i < matches.size(); i++) {
        auto p1_hat = kps_list[0][matches[i].queryIdx].pt;
        auto p2_hat = kps_list[1][matches[i].trainIdx].pt;
        Eigen::Vector3d p1;
        p1 << p1_hat.x, p1_hat.y, 1;
        Eigen::Vector2d p2;
        p2 << p2_hat.x, p2_hat.y;

        problem.AddResidualBlock(     // 向问题中添加误差项
        // 使用自动求导，模板参数：误差类型，输出维度，输入维度，维数要与前面struct中一致
        new ceres::AutoDiffCostFunction<ProjCost, 1, 7>(
            new ProjCost(p2, p1, K, T21)
        ),
        nullptr,            // 核函数，这里不使用，为空
        abc                 // 待估计参数
        );
    }

    // 配置求解器
    ceres::Solver::Options options;     // 这里有很多配置项可以填
    options.linear_solver_type = ceres::DENSE_NORMAL_CHOLESKY;  // 增量方程如何求解
    options.minimizer_progress_to_stdout = true;   // 输出到cout

    ceres::Solver::Summary summary;                // 优化信息
    ceres::Solve(options, &problem, &summary);  // 开始优化
    std::cout << summary.BriefReport() << std::endl;
    std::cout << "estimated a,b,c = ";
    for (auto a:abc) std::cout << a << " ";
    std::cout << std::endl;

}